package util

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	_ "github.com/mattn/go-sqlite3"
	"log"
	"sync"
	"time"
)

var (
	DBPool *Pool
)

func init() {
	pool, err := NewPool(Driver_Class, DB, 5, 10)
	if err != nil {
		log.Fatal("failed to create connection pool: ", err)
	}
	DBPool = pool
	//defer pool.Close()
	//初始化agent表
	sqlStmt := fmt.Sprintf(`CREATE TABLE %s (id INTEGER PRIMARY KEY, agentIp TEXT,agentMac TEXT,logInfo TEXT,createTime TEXT,state INTEGER);`, TableAgent)
	initDB(TableAgent, sqlStmt)
	//初始化ServerInfo
	sqlStmt = fmt.Sprintf(`CREATE TABLE %s (id INTEGER PRIMARY KEY, agentIp TEXT,agentMac TEXT,serverName TEXT,describe TEXT,logPath TEXT,createTime TEXT,state INTEGER);`, TableServerInfo)
	initDB(TableServerInfo, sqlStmt)
	//初始化logInfo表
	sqlStmt = fmt.Sprintf(`CREATE TABLE %s (id INTEGER PRIMARY KEY,agentMac TEXT,serverName TEXT, content TEXT,createTime LONG,logType INTEGER);`, TableLogContent)
	initDB(TableLogContent, sqlStmt)
}

// Pool 类型，用于封装 SQLite 数据库的连接池。Pool 类型包含以下字段：
//
// MaxIdle：空闲连接的最大数量。
// MaxActive：同时使用的最大连接数。
// mu：用于保护连接池的互斥锁。
// idleConns：空闲连接的列表。
// activeConns：正在使用的连接的列表。
// waitingConns：正在等待获取连接的 goroutine 的通道。
// 该类型的 NewPool 函数用于创建连接池，它接受四个参数：
//
// driverName：SQLite 驱动名称。
// dataSourceName：数据源名称。
// maxIdle：最大空闲连接数
type Pool struct {
	MaxIdle   int
	MaxActive int

	closeMu       sync.RWMutex
	closed        bool
	closeChannels []chan struct{}

	mu           sync.Mutex
	idleConns    []*sql.DB
	activeConns  []*sql.DB
	waitingConns chan struct{}
}

func NewPool(driverName, dataSourceName string, maxIdle, maxActive int) (*Pool, error) {
	pool := &Pool{
		MaxIdle:      maxIdle,
		MaxActive:    maxActive,
		waitingConns: make(chan struct{}, maxActive),
	}

	for i := 0; i < maxIdle; i++ {
		conn, err := sql.Open(driverName, dataSourceName)
		if err != nil {
			return nil, err
		}
		pool.idleConns = append(pool.idleConns, conn)
	}
	return pool, nil
}

func (p *Pool) getConn(ctx context.Context) (*sql.DB, error) {
	select {
	case p.waitingConns <- struct{}{}:
	default:
		return nil, errors.New("connection pool is full")
	}

	p.mu.Lock()
	defer p.mu.Unlock()

	for {
		if p.closed {
			return nil, ErrPoolClosed
		}

		n := len(p.idleConns)
		if n > 0 {
			conn := p.idleConns[n-1]
			p.idleConns = p.idleConns[:n-1]
			p.activeConns = append(p.activeConns, conn)
			return conn, nil
		}

		if len(p.activeConns) >= p.MaxActive {
			p.mu.Unlock()
			select {
			case <-ctx.Done():
				return nil, ctx.Err()
			case <-p.waitingConns:
			}
			p.mu.Lock()
			continue
		}

		conn, err := sql.Open(Driver_Class, DB)
		if err != nil {
			p.mu.Unlock()
			return nil, err
		}

		p.activeConns = append(p.activeConns, conn)
		return conn, nil
	}
}

func (p *Pool) PutConn(conn *sql.DB) error {
	p.mu.Lock()
	defer p.mu.Unlock()

	for i := range p.activeConns {
		if p.activeConns[i] == conn {
			p.activeConns = append(p.activeConns[:i], p.activeConns[i+1:]...)
			break
		}
	}

	if p.closed {
		conn.Close()
		return nil
	}

	p.idleConns = append(p.idleConns, conn)

	select {
	case <-p.waitingConns:
	default:
	}

	for _, ch := range p.closeChannels {
		close(ch)
	}
	p.closeChannels = nil

	return nil
}

func (p *Pool) Query(query string, args ...interface{}) (*sql.Rows, error) {
	conn, err := p.getConn(context.Background())
	if err != nil {
		return nil, err
	}
	defer p.PutConn(conn)

	return conn.Query(query, args...)
}

func (p *Pool) Exec(query string, args ...interface{}) (sql.Result, error) {
	conn, err := p.getConn(context.Background())
	if err != nil {
		return nil, err
	}
	defer p.PutConn(conn)
	return conn.Exec(query, args...)
}

var (
	ErrPoolClosed = errors.New("connection pool is closed")
)

func (p *Pool) Close() error {
	p.closeMu.Lock()
	defer p.closeMu.Unlock()

	if p.closed {
		return ErrPoolClosed
	}

	p.closed = true

	close(p.waitingConns)
	for _, conn := range p.idleConns {
		conn.Close()
	}
	for _, conn := range p.activeConns {
		conn.Close()
	}

	ch := make(chan struct{})
	p.closeChannels = append(p.closeChannels, ch)
	select {
	case <-ch:
	case <-time.After(time.Second * 5):
	}

	return nil
}

// 初始化
func initDB(tableName string, sql string) {
	// 检查表是否存在
	sqlStmt := fmt.Sprintf("SELECT name FROM sqlite_master WHERE type='table' AND name='%s';", tableName)
	rows, err := DBPool.Query(sqlStmt)
	if err != nil {
		log.Fatal(err)
	}
	// 如果表不存在，则创建它
	if !rows.Next() {
		//sqlStmt = fmt.Sprintf(`CREATE TABLE %s (id INTEGER PRIMARY KEY, agentIp TEXT,agentMac TEXT,logInfo TEXT);`, tableName)
		_, err = DBPool.Exec(sql)
		if err != nil {
			log.Printf("%q: %s\n", err, sql)
			return
		}
	}
	defer rows.Close()
}
